package com.dragon.executor;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 内存监控与代码执行的最终包装器。
 * <p>
 * 该类的职责是在一个独立的子进程中：
 * 1. 启动内存监控。
 * 2. 在一个独立的“工人”线程中执行用户代码，以防止主线程被阻塞。
 * 3. 严格执行时间限制。
 * 4. 捕获用户代码的 stdout 和 stderr。
 * 5. 捕获用户代码的运行时异常，并设置非零退出码。
 * 6. 无论执行成功、失败还是超时，都能按预定格式输出结果（用户输出/错误信息 + 分隔符 + 内存）。
 * </p>
 */
public class MemoryTrackingWrapper {

    private static final String SEPARATOR = "---MEMORY_SEPARATOR---";

    public static void main(String[] args) {
        // 至少需要主类名和超时时间
        if (args.length < 2) {
            System.err.println("用法: MemoryTrackingWrapper <真正的main类> <超时毫秒数> [参数...]");
            System.exit(1);
        }

        // 解析命令行参数
        String realMainClassName = args[0];
        long timeLimit = Long.parseLong(args[1]);
        String[] realArgs = Arrays.copyOfRange(args, 2, args.length);

        // 初始化内存监控
        final MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean();
        final AtomicLong peakMemory = new AtomicLong(0);
        ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

        // 重定向标准输出和标准错误流，以便捕获用户代码的所有输出
        PrintStream originalOut = System.out;
        ByteArrayOutputStream capturedOutStream = new ByteArrayOutputStream();
        System.setOut(new PrintStream(capturedOutStream));

        PrintStream originalErr = System.err;
        ByteArrayOutputStream capturedErrStream = new ByteArrayOutputStream();
        System.setErr(new PrintStream(capturedErrStream));

        // 以固定频率轮询并更新峰值内存使用情况
        scheduler.scheduleAtFixedRate(() -> {
            long usedHeapMemory = memoryMXBean.getHeapMemoryUsage().getUsed();
            peakMemory.getAndUpdate(x -> Math.max(x, usedHeapMemory));
        }, 0, 20, TimeUnit.MILLISECONDS);

        // 设置退出码，使用 AtomicInteger 以确保线程安全
        // 默认值为 0 (成功)，如果用户代码出错，工人线程会将其更新为 1
        final AtomicInteger exitCode = new AtomicInteger(0);

        // 创建并配置执行用户代码的工人线程
        Thread userThread = new Thread(() -> {
            try {
                // 使用反射调用用户代码的 main 方法
                Class<?> realMainClass = Class.forName(realMainClassName);
                Method mainMethod = realMainClass.getMethod("main", String[].class);
                mainMethod.invoke(null, (Object) realArgs);
            } catch (InvocationTargetException e) {
                // invoke 抛出的异常，其内部通常包装了用户代码的真实异常 (如除零、空指针等)
                Throwable cause = e.getCause();
                cause.printStackTrace(new PrintStream(capturedErrStream));
                // 将退出码设置为 1，表示程序运行时出错 (Runtime Error)
                exitCode.set(1);
            } catch (Exception e) {
                // 其他异常
                e.printStackTrace(new PrintStream(capturedErrStream));
                exitCode.set(1);
            }
        });

        boolean isTimeout = false;

        // 开始执行并监控工人线程
        try {
            userThread.start();

            userThread.join(timeLimit);

            // 检查工人线程是否仍在运行
            if (userThread.isAlive()) {

                isTimeout = true;
                // 强制停止失控的线程。
                userThread.stop();
            }
        } catch (InterruptedException e) {
            // 如果主线程在等待时被中断
            Thread.currentThread().interrupt();
            e.printStackTrace(new PrintStream(capturedErrStream));
            // 标记为错误
            exitCode.set(1);
        } finally {

            scheduler.shutdown();

            // 做最后一次内存检查，确保捕获到所有内存使用
            long finalMemory = memoryMXBean.getHeapMemoryUsage().getUsed();
            peakMemory.getAndUpdate(x -> Math.max(x, finalMemory));

            // 恢复原始的输出流
            System.setOut(originalOut);
            System.setErr(originalErr);

            // 打印用户输出或超时信息
            if (isTimeout) {
                System.out.print("Time Limit Exceeded");
            } else {
                System.out.print(capturedOutStream);
            }

            // 打印所有捕获的错误信息
            System.err.print(capturedErrStream);
            // 打印分隔符和峰值内存（单位KB）
            System.out.print(" " + SEPARATOR + " " + peakMemory.get() / 1024);
        }

        // 使用最终记录的退出码退出进程
        System.exit(exitCode.get());
    }
}